#pragma once
#include "topo.h"
#include <vector>
#include <cstdint>
#include <atomic>
#include <random>
#include <unordered_map>

struct PG {
    std::vector<std::pair<int, int>> disk_locs; // (server_idx, disk_idx)
};

struct ChunkMeta {
    uint64_t chunk_id;
    int pg_id;
};

class Pool {
public:
    bool init(const Topo& topo, int ec_data, int ec_parity);
    const std::vector<PG>& pgs() const { return pgs_; }
    uint64_t alloc_chunk_id();
    int get_ec_data() const { return ec_data_; }
    int get_ec_parity() const { return ec_parity_; }
    int get_pg_count() const { return pgs_.size(); }
    int allocated_chunk_count() const;
    std::vector<uint64_t> allocated_chunks() const;
    ChunkMeta alloc_chunkid_balanced();
    const ChunkMeta* get_chunk_meta(uint64_t chunkid) const;
    const PG* get_pg(int pgid) const;
    int get_pg_of_chunk(uint64_t chunkid) const;
private:
    std::vector<PG> pgs_;
    std::vector<std::pair<int, int>> all_disks_; // (server_idx, disk_idx)
    std::atomic<uint64_t> chunk_id_{1};
    int ec_data_ = 0;
    int ec_parity_ = 0;
    std::unordered_map<uint64_t, ChunkMeta> chunkid_meta_;
    std::vector<int> pg_chunk_count_;
}; 